# @file    : rag_llm_test
# @time    : 2025/4/3
# @author  : yongpeng.yao
# @desc    :
from pathlib import Path

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex

a_path = Path("data/电商A-Third Quarter 2023 Results.pdf")
b_path = Path("data/电商B-Third Quarter 2023 Results.pdf")

A_docs = SimpleDirectoryReader(input_files=[a_path]).load_data()
B_docs = SimpleDirectoryReader(input_files=[b_path]).load_data()

from deepseek_llm import DeepSeekLLM

llm = DeepSeekLLM()

A_index = VectorStoreIndex.from_documents(A_docs, llm=llm)
B_index = VectorStoreIndex.from_documents(B_docs, llm=llm)

A_index.storage_context.persist("storage/A_index")
B_index.storage_context.persist("storage/B_index")

from llama_index.core import StorageContext, load_index_from_storage

index_loaded = False

try:
    storage_context = StorageContext.from_defaults(persist_dir="./storge/A_index")
    A_index = load_index_from_storage(storage_context)
    storage_context = StorageContext.from_defaults(persist_dir="./storge/B_index")
    B_index = load_index_from_storage(storage_context)
    index_loaded = True
except Exception as e:
    _ = e
    index_loaded = False
print('index_loaded:', index_loaded)
a_engine = A_index.as_query_engine(similarity_top_k=3)
b_engine = B_index.as_query_engine(similarity_top_k=3)

from llama_index.core.tools import QueryEngineTool, ToolMetadata

query_engine_tools = [
    QueryEngineTool(
        query_engine=a_engine,
        metadata=ToolMetadata(name="A_Finance", description="用于提供电商A的财务信息"),
    ),
    QueryEngineTool(
        query_engine=b_engine,
        metadata=ToolMetadata(name="B_Finance", description="用于提供电商B的财务信息"),
    ),
]

from llama_index.core.agent import ReActAgent

agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True)

agent.chat('尽你所能用中文回答以下问题。比较一下两家电商的销售额')
